# recommend_next_day.py
import os
import pandas as pd
from typing import List, Optional, Tuple

import qlib
from qlib.workflow import R
from qlib.data import D


def _get_latest_recorder(experiment_name: str = "workflow"):
    exp = R.get_exp(experiment_name=experiment_name)
    recorders = exp.list_recorders()
    rid = sorted(recorders.keys())[-1]
    return exp, recorders[rid], rid


def _load_model_and_dataset(experiment_name: str = "workflow", recorder_id: Optional[str] = None):
    if recorder_id is None:
        _, recorder, _ = _get_latest_recorder(experiment_name)
    else:
        recorder = R.get_recorder(experiment_name=experiment_name, recorder_id=recorder_id)

    # 模型可能存成 "trained_model" 或 "params.pkl"
    try:
        model = recorder.load_object("trained_model")
    except Exception:
        model = recorder.load_object("params.pkl")

    dataset = recorder.load_object("dataset")
    return model, dataset, recorder


def _next_trading_day(from_date: pd.Timestamp) -> pd.Timestamp:
    # 从 from_date 开始找下一个交易日
    cal = D.calendar(start_time=from_date, end_time=from_date + pd.Timedelta(days=30), freq="day")
    cal = pd.DatetimeIndex(cal)
    
    # 找到第一个大于 from_date 的交易日
    next_trading_days = cal[cal > from_date]
    
    if len(next_trading_days) == 0:
        # 如果没有找到下一个交易日，说明数据可能不够新
        # 我们找最后一个可用的交易日
        if len(cal) > 0:
            last_available = cal[-1]
            print(f"Warning: No future trading days found. Using last available trading day: {last_available.date()}")
            return last_available
        else:
            raise ValueError(f"No trading days found in calendar after {from_date}")
    
    return next_trading_days[0]


def recommend_next_day(
    provider_uri: str,
    region: str = "cn",
    experiment_name: str = "workflow",
    recorder_id: Optional[str] = None,
    trade_date: Optional[str] = None,
    topk: int = 10,
    mlflow_uri: Optional[str] = None,
) -> Tuple[pd.Timestamp, List[str], pd.Series]:
    qlib.init(provider_uri=provider_uri, region=region, expression_cache=None, dataset_cache=None)

    if mlflow_uri:
        R.set_uri(mlflow_uri)

    model, dataset, recorder = _load_model_and_dataset(experiment_name=experiment_name, recorder_id=recorder_id)

    # 首先尝试直接加载已有的预测结果
    try:
        print("Trying to load existing prediction results...")
        pred_file = recorder.load_object("pred.pkl")
        print(f"Successfully loaded prediction file with shape: {pred_file.shape}")
        print(f"Prediction file type: {type(pred_file)}")
        print(f"Prediction file columns: {pred_file.columns if hasattr(pred_file, 'columns') else 'No columns'}")
        
        # 获取可用的日期
        available_dates = pred_file.index.get_level_values('datetime').unique()
        print(f"Available dates: {[d.date() for d in available_dates]}")
        
        if trade_date is None:
            # 使用最后一个可用日期
            trade_dt = available_dates[-1]
            print(f"Using last available date: {trade_dt.date()}")
        else:
            trade_dt = pd.Timestamp(trade_date)
            if trade_dt not in available_dates:
                print(f"Warning: {trade_dt.date()} not in available dates, using last available: {available_dates[-1].date()}")
                trade_dt = available_dates[-1]
        
        # 提取该日期的预测结果
        try:
            pred_day = pred_file.xs(trade_dt, level="datetime")
        except Exception as e:
            print(f"Failed to extract data for {trade_dt.date()}: {e}")
            # 如果 xs 失败，尝试其他方法
            if hasattr(pred_file, 'loc'):
                pred_day = pred_file.loc[pd.IndexSlice[:, trade_dt], :]
            else:
                raise
        
        # 确保 pred_day 是 Series 格式
        if hasattr(pred_day, 'columns') and len(pred_day.columns) > 0:
            # 如果是 DataFrame，取第一列作为分数
            pred_day = pred_day.iloc[:, 0]
        
        # 设置名称
        pred_day.name = "score"
        
        # 安全地排序
        if hasattr(pred_day, 'sort_values'):
            sorted_pred = pred_day.sort_values(ascending=False)
        else:
            # 如果不是 Series，尝试转换为 Series
            if hasattr(pred_day, 'values'):
                pred_day = pd.Series(pred_day.values, index=pred_day.index, name="score")
                sorted_pred = pred_day.sort_values(ascending=False)
            else:
                raise ValueError(f"Cannot sort prediction data of type: {type(pred_day)}")
        
        topk_instruments = sorted_pred.head(topk).index.tolist()
        return trade_dt, topk_instruments, sorted_pred
        
    except Exception as e:
        print(f"Failed to load existing predictions: {e}")
        print("This suggests the model was not fully trained or predictions were not saved.")
        raise ValueError(
            "Cannot load predictions. Please ensure:\n"
            "1. The model was fully trained with qrun\n"
            "2. The workflow included SignalRecord to generate predictions\n"
            "3. The mlflow_uri points to the correct training directory"
        )


if __name__ == "__main__":
    # 1) 修改为你的数据目录（例如你本机的 qlib 数据目录）
    #    常见为 ~/.qlib/qlib_data/cn_data 或自建的数据路径
    provider_uri = os.path.expanduser("~/.qlib/qlib_data/cn_data")

    # 2) 可选：指定 recorder_id；不指定则默认取最新一次 qrun 训练的记录
    recorder_id = None  # 如需指定，例如 "20250101-153000"

    # 3) 目标交易日：不填则自动取“今天之后的下一个交易日”
    trade_date = None  # 例如 "2025-01-10"

    trade_dt, buy_list, ranked_scores = recommend_next_day(
        provider_uri=provider_uri,
        region="cn",
        experiment_name="workflow",
        recorder_id=recorder_id,
        trade_date=trade_date,
        topk=10,
    )

    print(f"Trade date: {trade_dt.date()}")
    print("Buy candidates (Top-10):")
    for s in buy_list:
        print(s)
    # 如需查看完整排序：
    # print(ranked_scores.head(50))
